{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# 变换\n",
        "在本节中，将深入探讨 Relax 程序的变换。变换是编译流程中的关键组成部分，用于优化并与硬件后端进行集成。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "首先，按照在 [上一节](relax-creation) 中所做的那样，创建简单的 Relax 程序。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "collapsed": false,
        "tags": [
          "hide-cell"
        ]
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
              "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
              "<span style=\"color: #007979; font-style: italic\"># from tvm.script import relax as R</span>\n",
              "\n",
              "<span style=\"color: #AA22FF\">@I</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
              "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
              "    <span style=\"color: #AA22FF\">@R</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>function\n",
              "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">forward</span>(x: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #BA2121\">&quot;n&quot;</span>, <span style=\"color: #008000\">784</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), fc1_weight: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">784</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), fc1_bias: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">128</span>,), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), fc2_weight: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), fc2_bias: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">10</span>,), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #BA2121\">&quot;n&quot;</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>):\n",
              "        n <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64()\n",
              "        R<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;num_input&quot;</span>: <span style=\"color: #008000\">1</span>})\n",
              "        <span style=\"color: #008000; font-weight: bold\">with</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>dataflow():\n",
              "            permute_dims: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">784</span>, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>permute_dims(fc1_weight, axes<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">None</span>)\n",
              "            matmul: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((n, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>matmul(x, permute_dims, out_dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;void&quot;</span>)\n",
              "            add: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((n, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>add(matmul, fc1_bias)\n",
              "            relu: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((n, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(add)\n",
              "            permute_dims1: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>permute_dims(fc2_weight, axes<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">None</span>)\n",
              "            matmul1: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((n, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>matmul(relu, permute_dims1, out_dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;void&quot;</span>)\n",
              "            add1: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((n, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>add(matmul1, fc2_bias)\n",
              "            gv: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((n, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> add1\n",
              "            R<span style=\"color: #AA22FF; font-weight: bold\">.</span>output(gv)\n",
              "        <span style=\"color: #008000; font-weight: bold\">return</span> gv\n",
              "</pre></div>\n"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "import tvm\n",
        "from tvm import IRModule, relax\n",
        "from tvm.relax.frontend import nn\n",
        "\n",
        "\n",
        "class NNModule(nn.Module):\n",
        "    def __init__(self):\n",
        "        super().__init__()\n",
        "        self.fc1 = nn.Linear(784, 128)\n",
        "        self.relu1 = nn.ReLU()\n",
        "        self.fc2 = nn.Linear(128, 10)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.fc1(x)\n",
        "        x = self.relu1(x)\n",
        "        x = self.fc2(x)\n",
        "        return x\n",
        "\n",
        "\n",
        "origin_mod, params = NNModule().export_tvm(\n",
        "    {\"forward\": {\"x\": nn.spec.Tensor((\"n\", 784), \"float32\")}}\n",
        ")\n",
        "origin_mod.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 应用变换\n",
        "\n",
        "Pass 是对程序应用变换的主要方式。可以对程序应用 Pass。作为第一步，让应用内置的 {py:class}`~tvm.relax.transform.LegalizeOps` Pass，将高级算子降级为低级算子。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "collapsed": false,
        "tags": [
          "hide-output"
        ]
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
              "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
              "<span style=\"color: #007979; font-style: italic\"># from tvm.script import relax as R</span>\n",
              "\n",
              "<span style=\"color: #AA22FF\">@I</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
              "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
              "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func(private<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
              "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">add</span>(var_matmul: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle, fc1_bias: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>),), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), var_T_add: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle):\n",
              "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>bool(<span style=\"color: #008000; font-weight: bold\">True</span>)})\n",
              "        n <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64()\n",
              "        matmul <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(var_matmul, (n, T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>)))\n",
              "        T_add <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(var_T_add, (n, T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>)))\n",
              "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> ax0, ax1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(n, T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>)):\n",
              "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;T_add&quot;</span>):\n",
              "                v_ax0, v_ax1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SS&quot;</span>, [ax0, ax1])\n",
              "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(matmul[v_ax0, v_ax1], fc1_bias[v_ax1])\n",
              "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(T_add[v_ax0, v_ax1])\n",
              "                T_add[v_ax0, v_ax1] <span style=\"color: #AA22FF; font-weight: bold\">=</span> matmul[v_ax0, v_ax1] <span style=\"color: #AA22FF; font-weight: bold\">+</span> fc1_bias[v_ax1]\n",
              "\n",
              "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func(private<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
              "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">add1</span>(var_matmul1: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle, fc2_bias: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">10</span>),), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), var_T_add: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle):\n",
              "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>bool(<span style=\"color: #008000; font-weight: bold\">True</span>)})\n",
              "        n <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64()\n",
              "        matmul1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(var_matmul1, (n, T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">10</span>)))\n",
              "        T_add <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(var_T_add, (n, T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">10</span>)))\n",
              "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> ax0, ax1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(n, T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">10</span>)):\n",
              "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;T_add&quot;</span>):\n",
              "                v_ax0, v_ax1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SS&quot;</span>, [ax0, ax1])\n",
              "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(matmul1[v_ax0, v_ax1], fc2_bias[v_ax1])\n",
              "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(T_add[v_ax0, v_ax1])\n",
              "                T_add[v_ax0, v_ax1] <span style=\"color: #AA22FF; font-weight: bold\">=</span> matmul1[v_ax0, v_ax1] <span style=\"color: #AA22FF; font-weight: bold\">+</span> fc2_bias[v_ax1]\n",
              "\n",
              "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func(private<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
              "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">matmul</span>(var_x: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle, permute_dims: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">784</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), var_matmul: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle):\n",
              "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>bool(<span style=\"color: #008000; font-weight: bold\">True</span>)})\n",
              "        n <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64()\n",
              "        x <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(var_x, (n, T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">784</span>)))\n",
              "        matmul <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(var_matmul, (n, T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>)))\n",
              "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> i0, i1, k <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(n, T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">784</span>)):\n",
              "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;matmul&quot;</span>):\n",
              "                v_i0, v_i1, v_k <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SSR&quot;</span>, [i0, i1, k])\n",
              "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(x[v_i0, v_k], permute_dims[v_k, v_i1])\n",
              "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(matmul[v_i0, v_i1])\n",
              "                <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>init():\n",
              "                    matmul[v_i0, v_i1] <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">0.0</span>)\n",
              "                matmul[v_i0, v_i1] <span style=\"color: #AA22FF; font-weight: bold\">=</span> matmul[v_i0, v_i1] <span style=\"color: #AA22FF; font-weight: bold\">+</span> x[v_i0, v_k] <span style=\"color: #AA22FF; font-weight: bold\">*</span> permute_dims[v_k, v_i1]\n",
              "\n",
              "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func(private<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
              "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">matmul1</span>(var_relu: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle, permute_dims1: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">10</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), var_matmul: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle):\n",
              "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>bool(<span style=\"color: #008000; font-weight: bold\">True</span>)})\n",
              "        n <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64()\n",
              "        relu <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(var_relu, (n, T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>)))\n",
              "        matmul <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(var_matmul, (n, T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">10</span>)))\n",
              "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> i0, i1, k <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(n, T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">10</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>)):\n",
              "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;matmul&quot;</span>):\n",
              "                v_i0, v_i1, v_k <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SSR&quot;</span>, [i0, i1, k])\n",
              "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(relu[v_i0, v_k], permute_dims1[v_k, v_i1])\n",
              "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(matmul[v_i0, v_i1])\n",
              "                <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>init():\n",
              "                    matmul[v_i0, v_i1] <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">0.0</span>)\n",
              "                matmul[v_i0, v_i1] <span style=\"color: #AA22FF; font-weight: bold\">=</span> matmul[v_i0, v_i1] <span style=\"color: #AA22FF; font-weight: bold\">+</span> relu[v_i0, v_k] <span style=\"color: #AA22FF; font-weight: bold\">*</span> permute_dims1[v_k, v_i1]\n",
              "\n",
              "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func(private<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
              "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">relu</span>(var_add: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle, var_compute: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle):\n",
              "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>bool(<span style=\"color: #008000; font-weight: bold\">True</span>)})\n",
              "        n <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64()\n",
              "        add <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(var_add, (n, T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>)))\n",
              "        compute <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(var_compute, (n, T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>)))\n",
              "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> i0, i1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(n, T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>)):\n",
              "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;compute&quot;</span>):\n",
              "                v_i0, v_i1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SS&quot;</span>, [i0, i1])\n",
              "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(add[v_i0, v_i1])\n",
              "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(compute[v_i0, v_i1])\n",
              "                compute[v_i0, v_i1] <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>max(add[v_i0, v_i1], T<span style=\"color: #AA22FF; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">0.0</span>))\n",
              "\n",
              "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func(private<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
              "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">transpose</span>(fc1_weight: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">784</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), T_transpose: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">784</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>)):\n",
              "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>bool(<span style=\"color: #008000; font-weight: bold\">True</span>)})\n",
              "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> ax0, ax1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">784</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>)):\n",
              "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;T_transpose&quot;</span>):\n",
              "                v_ax0, v_ax1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SS&quot;</span>, [ax0, ax1])\n",
              "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(fc1_weight[v_ax1, v_ax0])\n",
              "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(T_transpose[v_ax0, v_ax1])\n",
              "                T_transpose[v_ax0, v_ax1] <span style=\"color: #AA22FF; font-weight: bold\">=</span> fc1_weight[v_ax1, v_ax0]\n",
              "\n",
              "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func(private<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
              "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">transpose1</span>(fc2_weight: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">10</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), T_transpose: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">10</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>)):\n",
              "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>bool(<span style=\"color: #008000; font-weight: bold\">True</span>)})\n",
              "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> ax0, ax1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">10</span>)):\n",
              "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;T_transpose&quot;</span>):\n",
              "                v_ax0, v_ax1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SS&quot;</span>, [ax0, ax1])\n",
              "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(fc2_weight[v_ax1, v_ax0])\n",
              "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(T_transpose[v_ax0, v_ax1])\n",
              "                T_transpose[v_ax0, v_ax1] <span style=\"color: #AA22FF; font-weight: bold\">=</span> fc2_weight[v_ax1, v_ax0]\n",
              "\n",
              "    <span style=\"color: #AA22FF\">@R</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>function\n",
              "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">forward</span>(x: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #BA2121\">&quot;n&quot;</span>, <span style=\"color: #008000\">784</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), fc1_weight: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">784</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), fc1_bias: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">128</span>,), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), fc2_weight: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), fc2_bias: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">10</span>,), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #BA2121\">&quot;n&quot;</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>):\n",
              "        n <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64()\n",
              "        R<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;num_input&quot;</span>: <span style=\"color: #008000\">1</span>})\n",
              "        cls <span style=\"color: #AA22FF; font-weight: bold\">=</span> Module\n",
              "        <span style=\"color: #008000; font-weight: bold\">with</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>dataflow():\n",
              "            permute_dims <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>call_tir(cls<span style=\"color: #AA22FF; font-weight: bold\">.</span>transpose, (fc1_weight,), out_sinfo<span style=\"color: #AA22FF; font-weight: bold\">=</span>R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">784</span>, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>))\n",
              "            matmul <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>call_tir(cls<span style=\"color: #AA22FF; font-weight: bold\">.</span>matmul, (x, permute_dims), out_sinfo<span style=\"color: #AA22FF; font-weight: bold\">=</span>R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((n, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>))\n",
              "            add <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>call_tir(cls<span style=\"color: #AA22FF; font-weight: bold\">.</span>add, (matmul, fc1_bias), out_sinfo<span style=\"color: #AA22FF; font-weight: bold\">=</span>R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((n, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>))\n",
              "            relu <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>call_tir(cls<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu, (add,), out_sinfo<span style=\"color: #AA22FF; font-weight: bold\">=</span>R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((n, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>))\n",
              "            permute_dims1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>call_tir(cls<span style=\"color: #AA22FF; font-weight: bold\">.</span>transpose1, (fc2_weight,), out_sinfo<span style=\"color: #AA22FF; font-weight: bold\">=</span>R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>))\n",
              "            matmul1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>call_tir(cls<span style=\"color: #AA22FF; font-weight: bold\">.</span>matmul1, (relu, permute_dims1), out_sinfo<span style=\"color: #AA22FF; font-weight: bold\">=</span>R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((n, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>))\n",
              "            add1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>call_tir(cls<span style=\"color: #AA22FF; font-weight: bold\">.</span>add1, (matmul1, fc2_bias), out_sinfo<span style=\"color: #AA22FF; font-weight: bold\">=</span>R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((n, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>))\n",
              "            gv: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((n, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> add1\n",
              "            R<span style=\"color: #AA22FF; font-weight: bold\">.</span>output(gv)\n",
              "        <span style=\"color: #008000; font-weight: bold\">return</span> gv\n",
              "</pre></div>\n"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "mod = tvm.relax.transform.LegalizeOps()(origin_mod)\n",
        "mod.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "从输出可以看到，程序中的高级算子（即 ``relax.op``）已被相应的低级算子（即 ``relax.call_tir``）所取代。\n",
        "\n",
        "接下来，尝试应用算子融合，这是机器学习编译器中广泛使用的一种优化技术。请注意，在 Relax 中，融合优化是通过一系列 Pass 的协作完成的。可以按顺序应用这些 Pass。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "collapsed": false,
        "tags": [
          "hide-output"
        ]
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
              "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
              "<span style=\"color: #007979; font-style: italic\"># from tvm.script import relax as R</span>\n",
              "\n",
              "<span style=\"color: #AA22FF\">@I</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
              "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
              "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func(private<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
              "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">fused_matmul1_add1</span>(p_relu: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle, permute_dims1: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">10</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), fc2_bias: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">10</span>),), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), p_output0: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle):\n",
              "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>bool(<span style=\"color: #008000; font-weight: bold\">True</span>)})\n",
              "        n <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64()\n",
              "        relu <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(p_relu, (n, T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>)))\n",
              "        T_add_intermediate <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(p_output0, (n, T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">10</span>)))\n",
              "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
              "        matmul_intermediate <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>alloc_buffer((n, T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">10</span>)))\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> i0, i1, k <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(n, T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">10</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>)):\n",
              "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;matmul&quot;</span>):\n",
              "                v_i0, v_i1, v_k <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SSR&quot;</span>, [i0, i1, k])\n",
              "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(relu[v_i0, v_k], permute_dims1[v_k, v_i1])\n",
              "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(matmul_intermediate[v_i0, v_i1])\n",
              "                <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>init():\n",
              "                    matmul_intermediate[v_i0, v_i1] <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">0.0</span>)\n",
              "                matmul_intermediate[v_i0, v_i1] <span style=\"color: #AA22FF; font-weight: bold\">=</span> matmul_intermediate[v_i0, v_i1] <span style=\"color: #AA22FF; font-weight: bold\">+</span> relu[v_i0, v_k] <span style=\"color: #AA22FF; font-weight: bold\">*</span> permute_dims1[v_k, v_i1]\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> ax0, ax1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(n, T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">10</span>)):\n",
              "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;T_add&quot;</span>):\n",
              "                v_ax0, v_ax1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SS&quot;</span>, [ax0, ax1])\n",
              "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(matmul_intermediate[v_ax0, v_ax1], fc2_bias[v_ax1])\n",
              "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(T_add_intermediate[v_ax0, v_ax1])\n",
              "                T_add_intermediate[v_ax0, v_ax1] <span style=\"color: #AA22FF; font-weight: bold\">=</span> matmul_intermediate[v_ax0, v_ax1] <span style=\"color: #AA22FF; font-weight: bold\">+</span> fc2_bias[v_ax1]\n",
              "\n",
              "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func(private<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
              "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">fused_matmul_add_relu</span>(p_x: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle, permute_dims: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">784</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), fc1_bias: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>),), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), p_output0: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle):\n",
              "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>bool(<span style=\"color: #008000; font-weight: bold\">True</span>)})\n",
              "        n <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64()\n",
              "        x <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(p_x, (n, T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">784</span>)))\n",
              "        compute_intermediate <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(p_output0, (n, T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>)))\n",
              "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
              "        matmul_intermediate <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>alloc_buffer((n, T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>)))\n",
              "        T_add_intermediate <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>alloc_buffer((n, T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>)))\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> i0, i1, k <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(n, T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">784</span>)):\n",
              "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;matmul&quot;</span>):\n",
              "                v_i0, v_i1, v_k <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SSR&quot;</span>, [i0, i1, k])\n",
              "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(x[v_i0, v_k], permute_dims[v_k, v_i1])\n",
              "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(matmul_intermediate[v_i0, v_i1])\n",
              "                <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>init():\n",
              "                    matmul_intermediate[v_i0, v_i1] <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">0.0</span>)\n",
              "                matmul_intermediate[v_i0, v_i1] <span style=\"color: #AA22FF; font-weight: bold\">=</span> matmul_intermediate[v_i0, v_i1] <span style=\"color: #AA22FF; font-weight: bold\">+</span> x[v_i0, v_k] <span style=\"color: #AA22FF; font-weight: bold\">*</span> permute_dims[v_k, v_i1]\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> ax0, ax1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(n, T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>)):\n",
              "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;T_add&quot;</span>):\n",
              "                v_ax0, v_ax1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SS&quot;</span>, [ax0, ax1])\n",
              "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(matmul_intermediate[v_ax0, v_ax1], fc1_bias[v_ax1])\n",
              "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(T_add_intermediate[v_ax0, v_ax1])\n",
              "                T_add_intermediate[v_ax0, v_ax1] <span style=\"color: #AA22FF; font-weight: bold\">=</span> matmul_intermediate[v_ax0, v_ax1] <span style=\"color: #AA22FF; font-weight: bold\">+</span> fc1_bias[v_ax1]\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> i0, i1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(n, T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>)):\n",
              "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;compute&quot;</span>):\n",
              "                v_i0, v_i1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SS&quot;</span>, [i0, i1])\n",
              "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(T_add_intermediate[v_i0, v_i1])\n",
              "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(compute_intermediate[v_i0, v_i1])\n",
              "                compute_intermediate[v_i0, v_i1] <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>max(T_add_intermediate[v_i0, v_i1], T<span style=\"color: #AA22FF; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">0.0</span>))\n",
              "\n",
              "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func(private<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
              "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">transpose</span>(fc1_weight: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">784</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), T_transpose: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">784</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>)):\n",
              "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;op_pattern&quot;</span>: <span style=\"color: #008000\">2</span>, <span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>bool(<span style=\"color: #008000; font-weight: bold\">True</span>)})\n",
              "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> ax0, ax1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">784</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>)):\n",
              "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;T_transpose&quot;</span>):\n",
              "                v_ax0, v_ax1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SS&quot;</span>, [ax0, ax1])\n",
              "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(fc1_weight[v_ax1, v_ax0])\n",
              "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(T_transpose[v_ax0, v_ax1])\n",
              "                T_transpose[v_ax0, v_ax1] <span style=\"color: #AA22FF; font-weight: bold\">=</span> fc1_weight[v_ax1, v_ax0]\n",
              "\n",
              "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func(private<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
              "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">transpose1</span>(fc2_weight: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">10</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), T_transpose: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">10</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>)):\n",
              "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;op_pattern&quot;</span>: <span style=\"color: #008000\">2</span>, <span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>bool(<span style=\"color: #008000; font-weight: bold\">True</span>)})\n",
              "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> ax0, ax1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">10</span>)):\n",
              "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;T_transpose&quot;</span>):\n",
              "                v_ax0, v_ax1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SS&quot;</span>, [ax0, ax1])\n",
              "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(fc2_weight[v_ax1, v_ax0])\n",
              "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(T_transpose[v_ax0, v_ax1])\n",
              "                T_transpose[v_ax0, v_ax1] <span style=\"color: #AA22FF; font-weight: bold\">=</span> fc2_weight[v_ax1, v_ax0]\n",
              "\n",
              "    <span style=\"color: #AA22FF\">@R</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>function\n",
              "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">forward</span>(x: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #BA2121\">&quot;n&quot;</span>, <span style=\"color: #008000\">784</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), fc1_weight: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">784</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), fc1_bias: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">128</span>,), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), fc2_weight: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), fc2_bias: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">10</span>,), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #BA2121\">&quot;n&quot;</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>):\n",
              "        n <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64()\n",
              "        R<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;num_input&quot;</span>: <span style=\"color: #008000\">1</span>})\n",
              "        cls <span style=\"color: #AA22FF; font-weight: bold\">=</span> Module\n",
              "        <span style=\"color: #008000; font-weight: bold\">with</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>dataflow():\n",
              "            permute_dims <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>call_tir(cls<span style=\"color: #AA22FF; font-weight: bold\">.</span>transpose, (fc1_weight,), out_sinfo<span style=\"color: #AA22FF; font-weight: bold\">=</span>R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">784</span>, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>))\n",
              "            lv <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>call_tir(cls<span style=\"color: #AA22FF; font-weight: bold\">.</span>fused_matmul_add_relu, (x, permute_dims, fc1_bias), out_sinfo<span style=\"color: #AA22FF; font-weight: bold\">=</span>R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((n, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>))\n",
              "            permute_dims1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>call_tir(cls<span style=\"color: #AA22FF; font-weight: bold\">.</span>transpose1, (fc2_weight,), out_sinfo<span style=\"color: #AA22FF; font-weight: bold\">=</span>R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>))\n",
              "            gv <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>call_tir(cls<span style=\"color: #AA22FF; font-weight: bold\">.</span>fused_matmul1_add1, (lv, permute_dims1, fc2_bias), out_sinfo<span style=\"color: #AA22FF; font-weight: bold\">=</span>R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((n, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>))\n",
              "            R<span style=\"color: #AA22FF; font-weight: bold\">.</span>output(gv)\n",
              "        <span style=\"color: #008000; font-weight: bold\">return</span> gv\n",
              "</pre></div>\n"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "mod = tvm.ir.transform.Sequential(\n",
        "    [\n",
        "        tvm.relax.transform.AnnotateTIROpPattern(),\n",
        "        tvm.relax.transform.FuseOps(),\n",
        "        tvm.relax.transform.FuseTIR(),\n",
        "    ]\n",
        ")(mod)\n",
        "mod.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "结果显示，``matmul``、``add`` 和 ``relu`` 算子被融合到了内核中（即 ``call_tir``）。\n",
        "\n",
        "有关所有内置 Pass 的详细信息，请参阅 {py:mod}`tvm.relax.transform`。\n",
        "\n",
        "## 自定义 Pass\n",
        "\n",
        "也可以定义自己的 Pass。以下是将 ``relu`` 算子重写为 ``gelu`` 算子的示例。\n",
        "\n",
        "首先，需要编写 Relax IR Mutator 来执行重写。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from tvm.relax.expr_functor import PyExprMutator, mutator\n",
        "\n",
        "\n",
        "@mutator\n",
        "class ReluRewriter(PyExprMutator):\n",
        "    def __init__(self, mod):\n",
        "        super().__init__(mod)\n",
        "\n",
        "    def visit_call_(self, call: relax.Call) -> relax.Expr:\n",
        "        # visit the relax.Call expr, and only handle the case when op is relax.nn.relu\n",
        "        if call.op.name == \"relax.nn.relu\":\n",
        "            return relax.op.nn.gelu(call.args[0])\n",
        "\n",
        "        return super().visit_call_(call)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "然后，可以编写 Pass，将 Mutator 应用到整个模块中。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "collapsed": false,
        "tags": [
          "hide-output"
        ]
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
              "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
              "<span style=\"color: #007979; font-style: italic\"># from tvm.script import relax as R</span>\n",
              "\n",
              "<span style=\"color: #AA22FF\">@I</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
              "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
              "    <span style=\"color: #AA22FF\">@R</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>function\n",
              "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">forward</span>(x: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #BA2121\">&quot;n&quot;</span>, <span style=\"color: #008000\">784</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), fc1_weight: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">784</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), fc1_bias: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">128</span>,), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), fc2_weight: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), fc2_bias: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">10</span>,), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #BA2121\">&quot;n&quot;</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>):\n",
              "        n <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64()\n",
              "        R<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;num_input&quot;</span>: <span style=\"color: #008000\">1</span>})\n",
              "        <span style=\"color: #008000; font-weight: bold\">with</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>dataflow():\n",
              "            permute_dims: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">784</span>, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>permute_dims(fc1_weight, axes<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">None</span>)\n",
              "            matmul: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((n, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>matmul(x, permute_dims, out_dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;void&quot;</span>)\n",
              "            add: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((n, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>add(matmul, fc1_bias)\n",
              "            relu: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((n, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>gelu(add)\n",
              "            permute_dims1: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>permute_dims(fc2_weight, axes<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">None</span>)\n",
              "            matmul1: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((n, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>matmul(relu, permute_dims1, out_dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;void&quot;</span>)\n",
              "            add1: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((n, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>add(matmul1, fc2_bias)\n",
              "            gv: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((n, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> add1\n",
              "            R<span style=\"color: #AA22FF; font-weight: bold\">.</span>output(gv)\n",
              "        <span style=\"color: #008000; font-weight: bold\">return</span> gv\n",
              "</pre></div>\n"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "@tvm.transform.module_pass(opt_level=0, name=\"ReluToGelu\")\n",
        "class ReluToGelu:  # pylint: disable=too-few-public-methods\n",
        "    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:\n",
        "        \"\"\"IRModule-level transformation\"\"\"\n",
        "        rewriter = ReluRewriter(mod)\n",
        "        for g_var, func in mod.functions_items():\n",
        "            if isinstance(func, relax.Function):\n",
        "                func = rewriter.visit_expr(func)\n",
        "                rewriter.builder_.update_func(g_var, func)\n",
        "        return rewriter.builder_.get()\n",
        "\n",
        "\n",
        "mod = ReluToGelu()(origin_mod)\n",
        "mod.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "打印输出显示，``relax.nn.relu`` 运算符已被重写为 ``relax.nn.gelu`` 运算符。\n",
        "\n",
        "有关 Mutator 的详细信息，请参阅 {py:class}`~tvm.relax.expr_functor.PyExprMutator`。\n",
        "\n",
        "## 总结\n",
        "\n",
        "在本节中，展示了如何对 Relax 程序应用转换。还展示了如何定义和应用自定义转换。"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "ai",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.12.7"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
